package com.cuz.daileetcode;

import com.cuz.daileetcode.utils.graph.Edge;
import com.cuz.daileetcode.utils.graph.Graph;
import com.cuz.daileetcode.utils.graph.Node;

import java.util.*;

public class Day34 {


    public static class SimulationUnionSet<T> {
        public HashMap<Node<T>, List<Node<T>>> memory = new HashMap();

        /***
         * 根据图生成模拟并查集 将每个节点初始化指向自己的单个集合
         * @param graph 图
         */
        public static <T> SimulationUnionSet<T> create(Graph<T> graph) {
            SimulationUnionSet<T> res = new SimulationUnionSet<T>();
            for (Node<T> tempNode : Objects.requireNonNull(graph).nodes.values()) {
                res.memory.put(tempNode, Collections.singletonList(tempNode));
            }
            return res;
        }

        /**
         * 两个节点是否属于同一个集
         * @return
         */
        public boolean isSame(Node<T> node1, Node<T> node2) {
            return memory.get(node1) == memory.get(node2);
        }

        /**
         * 合并
         * 让前者的集合包括后者集合中的每一个元素
         * 并且让后者集合中的每一个节点指向前者集合
         */
        public void union(Node<T> node1, Node<T> node2) {
            if (isSame(node1, node2)) {
                return;
            }
            List<Node<T>> node1List = memory.get(node1);
            List<Node<T>> node2List = memory.get(node2);
            Objects.requireNonNull(node1List);
            Objects.requireNonNull(node2);
            for (Node<T> temNode : node2List) {
                node1List.add(temNode);
                memory.put(temNode, node1List);
            }
        }
    }

    /***
     * kruskal 生成最小生成树
     * @param graph
     * @param <T>
     * @return
     */
    public static <T> Set<Edge<T>> kruskal(Graph<T> graph) {
        Objects.requireNonNull(graph);
        //生成模拟并查集
        SimulationUnionSet<T> unionSet = SimulationUnionSet.create(graph);
        //将所有Edge 按照weight排序
        PriorityQueue<Edge<T>> heap = new PriorityQueue<>(Edge.edgeWeightComparator());
        //加入所有edge
        heap.addAll(graph.edges);
        //结果
        Set<Edge<T>> res = new HashSet<>();
        while (!heap.isEmpty()) {
            //拿到当前权重最小的边
            Edge<T> tempEdge = heap.poll();
            Node<T> from = tempEdge.from;
            Node<T> to = tempEdge.to;
            //两个节点是否属于同一个集，如果使代表选择会成环 那么继续下一轮
            //反之加入边，合并两个节点的集合
            if (!unionSet.isSame(from, to)) {
                unionSet.union(from,to);
                res.add(tempEdge);
            }
        }
        return res;
    }

}
